from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
from dolfinx import mesh, fem
from dolfinx.mesh import locate_entities_boundary, meshtags
from dolfinx.fem.petsc import assemble_matrix, assemble_vector
import ufl
import config as cfg

# ----------------- 基本设置 -----------------
problem_type = "smooth1"  # "smooth1", "smooth2", "singular"
type_ = "SIPG"            # "SIPG" | "NIPG" | "IIPG"
k = 2                     # 多项式阶数

# 物性与反应项（可按需修改）
K_val = 1.0               # K = K_val * I（标量各向同性）
alpha_val = 0.0           # 反应项系数 α

# 区域与精确解/派生量
if problem_type == "smooth1":
    x_min, x_max, y_min, y_max = 0.0, 1.0, 0.0, 1.0
    beta0 = 1.0  # 重要：惩罚项按 h^{-beta0} 缩放，一般取 1
    p_solution    = cfg.smooth_solution1
    p_solution_dx = cfg.smooth_solution1_dx
    p_solution_dy = cfg.smooth_solution1_dy
    
elif problem_type == "smooth2":
    x_min, x_max, y_min, y_max = 0.0, 1.0, 0.0, 1.0
    beta0 = 1.0
    p_solution    = cfg.smooth_solution2
    p_solution_dx = cfg.smooth_solution2_dx
    p_solution_dy = cfg.smooth_solution2_dy
elif problem_type == "singular":
    x_min, x_max, y_min, y_max = -1.0, 1.0, -1.0, 1.0
    beta0 = 1.0
    p_solution    = cfg.singular_solution
    p_solution_dx = cfg.singular_solution_dx
    p_solution_dy = cfg.singular_solution_dy
elif problem_type == "test":
    x_min, x_max, y_min, y_max = 0.0, 1.0, 0.0, 1.0
    beta0 = 1.0
    p_solution    = cfg.test_function
    p_solution_dx = cfg.test_function_dx
    p_solution_dy = cfg.test_function_dy

# IPG 族参数
if type_ == "NIPG":
    epsilon = +1
    # 重要：σ_0 的“基准常数”，后续会乘以 p*(p+1) 并除以 h
    sigma_0_base = 4.0
    sigma_1 = 0.0
elif type_ == "IIPG":
    epsilon = 0
    sigma_0_base = 8.0
    sigma_1 = 0.0
elif type_ == "SIPG":
    epsilon = -1
    sigma_0_base = 20  # 重要：配合 p*(p+1) 的缩放，保证强制性
    sigma_1 = 0.0

# 网格尺寸列表
h_list = [1/4, 1/8, 1/16, 1/32, 1/64]

# 误差收敛统计
h1_semi_err, l2_err = [], []

for h in h_list:
    N = int(round(1/h))
    domain = mesh.create_rectangle(
        MPI.COMM_WORLD,
        [[x_min, y_min], [x_max, y_max]],
        [N, N],
        mesh.CellType.triangle,
    )
    x = ufl.SpatialCoordinate(domain)

    # 有限元空间（DG-k）
    V = fem.functionspace(domain, ("DG", k))

    # Trial/Test
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)

    # 几何量与测度
    n = ufl.FacetNormal(domain)
    h_cell = ufl.CellDiameter(domain)   # 单元直径 这涉及两个单元 
    h_avg = ufl.avg(h_cell)             # 内部面平均 h
    
    sigma_int = sigma_0_base * (k + 1) * k / h_avg   # 内部面惩罚
    sigma_bnd = sigma_0_base * (k + 1) * k / h_cell     # 边界面惩罚

    dx = ufl.dx(domain)
    dS = ufl.dS(domain=domain)          # 内部面
    ds = ufl.ds(domain=domain)          # 外边界（稍后重定义用于标签）
    
    # np.concatenate用来把多个数组（array-like 对象）沿着指定的轴（axis）连接在一起。
    # 例如
    # a = np.array([1, 2, 3])
    # b = np.array([4, 5, 6])
    # c = np.array([7, 8, 9])
    # d = np.concatenate((a, b, c))
    # d = array([1, 2, 3, 4, 5, 6, 7, 8, 9])

    # 万恶之源
    # meshtags函数给网格实体打标签
    # 要点是entities必须严格升序排序 不能有重复索引
    # 错误原因是locate_entities_boundary(...) 分别对 left/right/top/bottom 做查询
    # 返回的本地 facet 索引各自有序，但跨集合拼接后整体未必有序。
    # np.concatenate([facets_left, facets_right, facets_bottom, facets_top])这里失去了有序性
    
    # 这里提供几个思路
    # 1. 反正是纯Dirichlet 干脆直接用ds进行标记
    # 2. 使用这种写法 但四条边分成四个标签t1~t4 积分时手动累加 《不推荐》
    # 3. 把四个边的facet索引合并后排序 并相应调整标签顺序 《不推荐》
    
    # 四个边
    # ---- 纯 Dirichlet 边界标记：把四条边都标成 Γ_D ----
    # tD = 1  # Dirichlet 标签 id

    # def on_left(x):   return np.isclose(x[0], x_min)
    # def on_right(x):  return np.isclose(x[0], x_max)
    # def on_bottom(x): return np.isclose(x[1], y_min)
    # def on_top(x):    return np.isclose(x[1], y_max)

    # facets_left   = locate_entities_boundary(domain, domain.topology.dim - 1, on_left)
    # facets_right  = locate_entities_boundary(domain, domain.topology.dim - 1, on_right)
    # facets_bottom = locate_entities_boundary(domain, domain.topology.dim - 1, on_bottom)
    # facets_top    = locate_entities_boundary(domain, domain.topology.dim - 1, on_top)

    # facet_indices = np.concatenate([facets_left, facets_right, facets_bottom, facets_top])
    # facet_tags    = np.full(facet_indices.size, tD, dtype=np.int32)

    # mt_facets = meshtags(domain, domain.topology.dim - 1, facet_indices, facet_tags)
    # # 重要：带 subdomain_data 的 ds，用 ds(tD) 选中纯 Dirichlet 边
    # ds = ufl.Measure("ds", domain=domain, subdomain_data=mt_facets)

    # ---- 系数/常量 ----
    K = K_val  # 标量各向同性
    alpha = fem.Constant(domain, PETSc.ScalarType(alpha_val))

    # ---- 制造解与边界/源项 ----
    if problem_type == "smooth1":
        u_exact = ufl.exp(-x[0] - x[1]*x[1])
        f = (1.0 - 4.0*x[1]*x[1]) * u_exact

    ux = fem.Function(V); ux.interpolate(lambda x: p_solution_dx(x[0], x[1]))
    uy = fem.Function(V); uy.interpolate(lambda x: p_solution_dy(x[0], x[1]))



    # ---- 惩罚参数（重要）：σ = C * p*(p+1) ----
    sigma0 = sigma_0_base * k * (k + 1)

    # ---- 双线性形式 a_ε(·,·) ----
    # 体项
    a_vol = K * ufl.dot(ufl.grad(u), ufl.grad(v)) * dx + alpha * u * v * dx

    # 重要：内部面项 —— 使用向量·向量的 inner(avg(grad(u)), jump(v))
    # a_int = - {K∇u} : [[v]] + ε {K∇v} : [[u]] + (σ/h) [[u]]:[[v]]
    a_int = (
        - ufl.inner(ufl.avg(K * ufl.grad(v)), ufl.jump(u, n)) * dS
        + epsilon * ufl.inner(ufl.avg(K * ufl.grad(u)), ufl.jump(v, n)) * dS
        + sigma_int * ufl.inner(ufl.jump(u, n), ufl.jump(v, n)) * dS
    )

    # 重要：边界（纯 Dirichlet）的 Nitsche 施加：与内部面一致的“镜像”结构
    # LHS: - (K∇u·n) v - ε (K∇v·n) u + (σ/h) u v
    a_D = (
        epsilon * K * ufl.dot(ufl.grad(u), n) * v * ds  # 符号修正
        - K * ufl.dot(ufl.grad(v), n) * u * ds
        + sigma_bnd * u * v * ds
    )

    a = a_vol + a_int + a_D

    # ---- 线性泛函 L(v) ----
    # L = (f, v) - ε (K∇v·n, g)_{Γ_D} + (σ/h) (g, v)_{Γ_D}
    L_vol = ufl.inner(f, v) * dx
    gD = u_exact  # 重要：纯 Dirichlet 的边界数据
    L_D = (
        + epsilon * K * ufl.dot(ufl.grad(v), n) * gD * ds
        + sigma_bnd * gD * v * ds
    )
    L = L_vol + L_D

    # ---- 组装并求解 ----
    a_form = fem.form(a)
    L_form = fem.form(L)

    A = assemble_matrix(a_form)
    A.assemble()
    b = assemble_vector(L_form)

    # 线性求解
    ksp = PETSc.KSP().create(MPI.COMM_WORLD)
    ksp.setOperators(A)
    ksp.setType("preonly")
    ksp.getPC().setType("lu")

    x_sol = A.createVecRight()
    uh = fem.Function(V)

    ksp.solve(b, x_sol)
    uh.x.array[:] = x_sol.array
    # print(uh.x.array - u_exact.x.array)
    

    # ---- 误差（broken H1-semi & L2）----
    l2_sq = fem.assemble_scalar(fem.form(((uh - u_exact)**2) * dx))
    l2_sq = MPI.COMM_WORLD.allreduce(l2_sq, op=MPI.SUM)
    l2_err.append(np.sqrt(l2_sq))

    grad_diff = ufl.grad(uh - u_exact)
    h1_sq = fem.assemble_scalar(fem.form(ufl.inner(grad_diff, grad_diff) * dx))
    h1_sq = MPI.COMM_WORLD.allreduce(h1_sq, op=MPI.SUM)
    h1_semi_err.append(np.sqrt(h1_sq))

    if MPI.COMM_WORLD.rank == 0:
        print(f"h = {h:.5f}:\tL2 err = {l2_err[-1]:.3e},\tH1-semi err = {h1_semi_err[-1]:.3e}")

# -------- 收敛阶统计 --------
def _rates(errs, hs):
    out = []
    for i in range(len(errs) - 1):
        out.append(np.log(errs[i] / errs[i+1]) / np.log(hs[i] / hs[i+1]))
    return out

l2_list = _rates(l2_err, h_list) if len(l2_err) > 1 else []
h1_semi_list = _rates(h1_semi_err, h_list) if len(h1_semi_err) > 1 else []

if MPI.COMM_WORLD.rank == 0:
    print("\n===== 收敛结果 =====")
    print("h:               ", ["{:.5f}".format(h) for h in h_list])
    if l2_list:
        print("L2收敛阶:         ", ["{:.3f}".format(e) for e in l2_list])
    if h1_semi_list:
        print("H1半范数收敛阶:   ", ["{:.3f}".format(e) for e in h1_semi_list])